from pydantic import BaseModel, Field
from app.bailian.common import chat_prompt_template, llm
from langchain_core.tools import tool

# 定义一个简单的加法工具并使用装饰器表明函数功能
class AddInputArgs(BaseModel):
    a: int = Field(description="The first number to add.")
    b: int = Field(description="The second number to add.")

@tool(
    description="Add two numbers together.",
    args_schema=AddInputArgs,
)
def add(a, b):
    return a + b

# 将工具函数注册到字典中以便后续调用
tool_dict = {
    "add": add,
}

# 将工具绑定到语言模型
llm_with_tools = llm.bind_tools(tool_dict.values())

# 使用聊天提示词模板和语言模型创建一个链
chain = chat_prompt_template | llm_with_tools
resp = chain.invoke({"role": "计算", "domain": "数学计算", "question": "使用工具计算：100+100等于多少？"})

# 处理工具调用结果
for tool_call in resp.tool_calls:
    func_name = tool_call["name"]
    func_args = tool_call["args"]
    tool_func = tool_dict[func_name]

    tool_content = tool_func.invoke(func_args)
    print("Final Response:", tool_content, end="", flush=True)
